import argparse  # 导入参数库
import datetime
import queue  # 导入队列库
import socket
import threading  # 导入线程库

import prettytable  # 导入表格库
# 作业1：argarps库指定参数做多线程扫描，需要用户可以手动指定线程和常用端口或者全端口（需要根据用户的输入调用不同的函数），将结果保存到文件txt中
# 作业2：写一个双写或大小写的绕过tamper
import pyfiglet  # 导入字体库
from colorama import Fore  # 导入颜色库


# 横幅
def banner():
    print(Fore.BLUE+"PortScan v1.0  简单的端口扫描工具   Author: Youze"  )
    test = pyfiglet.figlet_format("PortScan", font="calvin_s")
    print(Fore.BLUE+test)

# 命令参数获取
def get_args():
    parser = argparse.ArgumentParser(description='端口扫描工具')
    parser.add_argument('-u', '--target', help='目标地址', required=True) #required=True表示必须输入
    parser.add_argument('-v', '--version', action='version', version='%(prog)s 1.0', help='显示版本信息')
    #%(prog)s 1.0版本中，%(prog)s 被替换成了PortScan  action='version'表示显示版本信息
    parser.add_argument('-p', '--port', help='扫描端口范围 all 或 normal 或 start-end 或 port1,port2,port3', default='normal')
    parser.add_argument('-o', '--output', help='输出文件', default='')
    parser.add_argument('-n', '--number', help='线程数量', default=10, type=int)
    return parser.parse_args()

# 端口格式检查
def port_check(port):
    if port == 'normal':
        default_ports = [20, 21, 22, 23, 25, 53, 80, 110, 143, 443, 445, 993, 995, 1723, 3306, 3389, 5432, 8080, 8443, 1521, 5900, 6379, 10000, 33000]
        return default_ports
    elif port == 'all':
        return list(range(1, 65535+1))
    elif '-' in port:
        start_port, end_port = map(int, port.split('-'))
        if start_port < 1 or start_port > 65535 or end_port < 1 or end_port > 65535:
            print(Fore.RED+'端口范围错误')
            exit()
        elif start_port > end_port:
            temp = start_port
            start_port = end_port
            end_port = temp
        return list(range(start_port, end_port+1))
    else:
        return [int(p) for p in port.split(',')]

queue = queue.Queue()  # 队列
threads = []  # 线程列表
table = prettytable.PrettyTable(['Port', 'time'])  # 表格
def scan_port(host,queue):
    print(queue.empty())
    while not queue.empty():  # 队列不为空时循环
        port = queue.get()
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        result = sock.connect_ex((host, port))
        time = datetime.datetime.now()  #当前时间
        if result == 0:
            return [port, time]

def thread_scan(srgs_host, srgs_port, srgs_output, srgs_number):
    # 将srgs_port中的端口放入队列
    for p in srgs_port:
        queue.put(p)
    # 创建线程
    for i in range(srgs_number):
        t = threading.Thread(target=scan_port, args=(srgs_host,queue))
        threads.append(t)
        t.start()

    # 输出结果
    print(table)









if __name__ == '__main__':
    banner()
    srgs = get_args()
    srgs_host = srgs.target
    srgs_port = port_check(srgs.port)
    print(Fore.YELLOW+'[+]扫描端口：', srgs_port)
    srgs_output = srgs.output
    srgs_number = srgs.number
    print(Fore.YELLOW+'[+]开始扫描')
    port_result = thread_scan(srgs_host, srgs_port, srgs_output, srgs_number)








